from pprint import pprint
from typing import Literal, TypedDict
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.graph import StateGraph, START, END,MessagesState
from langgraph.graph.message import add_messages
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langgraph.prebuilt import ToolNode
from langgraph.checkpoint.memory import MemorySaver

from config.model_config import get_chat_openai_zhipu


@tool
def search(city: str) -> str:
    """获取指定城市的天气信息"""
    if "上海" in city.lower() or 'shanghai' in city.lower():
        # return '上海是中国的首都'
        return "现在30度, 有雾"
    return "现在35度, 阳光明媚"

# 工具类别
tools = [search]

# 创建工具节点
tool_nodes = ToolNode(tools)


# 实际请以官方文档为准
model = get_chat_openai_zhipu().bind_tools(tools)

def should_continue(state: MessagesState) -> Literal["tools", END]:
    """Check if the last message is a tool call."""
    messages = state["messages"]
    last_message = messages[-1]
    if last_message.tool_calls:
        return "tools"
    return END
    # if isinstance(state["messages"][-1], ToolMessage):
    #     return "tools"
    # return END

# 1.定义调用模型的函数
def call_model(state: MessagesState) -> MessagesState:
    """Call the model with the messages."""
    messages = state["messages"]
    # messages[HumanMessage(content='上海的天气怎么样', additional_kwargs={}, response_metadata={},
    #                       id='541dac4c-be5a-4391-8f16-d34d95c6409b')]
    # 调用模型
    # 提取消息内容，创建新的 HumanMessage 实例，不包含 id
    response = model.invoke(messages)
    # 将模型的响应添加到消息列表中
    # state["messages"].append(response)
    return {"messages": [response]}

# 2.用状态初始化，定义一个新的状态图
workflow = StateGraph(MessagesState)

# 3.定义图节点, 定义我们将循环的2个节点
workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_nodes)

# 4.定义入口点和图边
workflow.set_entry_point("agent")

# 添加条件边
workflow.add_conditional_edges(
    "agent",
    should_continue
)

# 定义图边
workflow.add_edge("tools", "agent")

#初始化内存在图运行之间持久化状态
checkpointer = MemorySaver() #可以存redis

# 5.编译
app = workflow.compile(checkpointer=checkpointer)

# # 6.执行图, 使用可运行对象
# final_response = app.invoke({"messages": [HumanMessage(content="上海的天气怎么样")]},
#                             config={"configurable": {"thread_id": "42"}})
#
# result = final_response['messages'][-1].content
# print(result)
# final_response = app.invoke({"messages": [HumanMessage(content="我问的哪个城市")]},
#                             config={"configurable": {"thread_id": "42"}})
# result = final_response['messages'][-1].content
# print(result)


from langchain_core.messages import AIMessageChunk, ToolMessage
while True:
    user_input = input("User: ")
    if user_input.lower() in ["quit", "exit", "q"]:
        print("Goodbye!")
        break
    for event in app.stream({"messages": [HumanMessage(content=user_input)]},
                                config={"configurable": {"thread_id": "42"}}, stream_mode="messages"):
        # print(event)
        if type(event[0]) == ToolMessage:
            print('调用了工具')
        elif type(event[0]) == AIMessageChunk:
            if event[0].content:
                print(event[0].content, end="")
    print('\n')